#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@Author  :   FRzhuizhu
@License :   (C) Copyright 2013-2020, FRzhuizhu
@Contact :   736019332@qq.com
@Software:   PyCharm
@File    :   heterostructure.py
@Time    :   2020/9/23 19:29
@Desc    :
'''

from cmath import pi, cos, sin, acos
import numpy as np
import time
from multiprocessing import Pool




class Vec:

    def __init__(self, data: list):
        self.raw = data

    def __add__(self, other):
        return Vec([self.raw[i]+other.raw[i] for i in range(len(self))])

    def __mul__(self, other):
        return Vec([other*x for x in self.raw])

    def __len__(self):
        return len(self.raw)

    def __repr__(self):
        return f'<Vec({self.raw})>'

def norm(a):
    return (sum([x**2 for x in a.raw]))**0.5


def dot(a, b):
    return sum([a.raw[i]*b.raw[i] for i in range(len(a))])



def get_form(aa, bb, ang, vv) -> np.ndarray:
    xit = ang * pi / 180.0
    a = Vec([cos(xit / 2.0) * aa, sin(xit / 2.0) * aa])
    b = Vec([cos(xit / 2.0) * bb, sin(xit / 2.0) * bb * (-1.0)])
    ss = np.zeros(7)

    for m1 in range(-vv, vv+1):
        for n1 in range(-vv, vv+1):
            for m2 in range(-vv, vv+1):
                for n2 in range(-vv, vv+1):

                    A = (a * m1) + (b * n1)
                    B = (a * m2) + (b * n2)
                    try:
                        db = dot(A, B)
                        temp_a = ((acos(dot(A, B) / (norm(A) * norm(B)))) * 180 / pi).real
                    except ZeroDivisionError:
                        continue

                    temp = np.zeros(7)
                    temp[0] = m1
                    temp[1] = n1
                    temp[2] = m2
                    temp[3] = n2
                    temp[4] = norm(A).real
                    temp[5] = norm(B).real
                    temp[6] = temp_a
                    ss = np.vstack((ss, temp))

                    # s1[k, 1] = m1
                    # s1[k, 2] = n1
                    # s1[k, 3] = m2
                    # s1[k, 4] = n2
                    # s1[k, 5] = norm(A)
                    # s1[k, 6] = norm(B)
    return ss[1:]


def get_s3(s1, s2, m, n):
    e1, e2 = s1[:, 6], s2[:, 6]
    t1, t2 = s1[:, 4], s2[:, 4]
    y1, y2 = s1[:, 5], s2[:, 5]
    # k = k + 1
    temp = np.zeros(15)
    # s3(k, 1: 4) = l1(m, 1: 4);
    temp[0:3] = s1[m, 0:3]
    # s3(k, 5: 8) = l2(n, 1: 4);
    temp[4:7] = s2[n, 0:3]
    # s3(k, 9) = e1(m, 1) - e2(n, 1);
    temp[8] = e1[m] - e2[n]
    # s3(k, 10) = e1(m, 1);
    temp[9] = e1[m]
    # s3(k, 11) = e2(n, 1);
    temp[10] = e2[n]
    # sss1 = t1(m, 1) - t2(n, 1);
    sss1 = t1[m] - t2[n]
    # sss2 = y1(m, 1) - y2(n, 1);
    sss2 = y1[m] - y2[n]
    # s3(k, 12) = sss1. / t1(m, 1);
    temp[11] = sss1 / t1[m]
    # s3(k, 13) = sss1. / t2(n, 1);
    temp[12] = sss1 / t2[n]
    # s3(k, 14) = sss2. / y1(m, 1);
    temp[13] = sss2 / y1[m]
    # s3(k, 15) = sss2. / y2(n, 1);
    temp[14] = sss2 / y2[n]
    if (m*len(s2)+n) % 1000 == 0:
        print(m*len(s2)+n)
    return temp


def get_s4(s3):
    mis_ang = 3
    mis_rate = 0.05
    qw1 = abs(s3[:, 8])
    qw2 = abs(s3[:, 11])
    qw3 = abs(s3[:, 12])
    qw4 = abs(s3[:, 13])
    qw5 = abs(s3[:, 14])
    index = np.where((qw1 < mis_ang) & (qw2 < mis_rate) & (qw3 < mis_rate) & (qw4 < mis_rate) & (qw5 < mis_rate))
    return s3[index]


def main():
    print('By FRzhuizhu')
    v = 2

    a1 = 3.776
    b1 = 9.486
    xita = 90
    N1 = 36

    a2 = 4.78069
    b2 = 4.78069
    beta = 120
    N2 = 7

    mis_ang = 3
    mis_rate = 0.05

    s1 = get_form(a1, b1, xita, v)
    np.savetxt('s1.csv', s1, fmt='%.5f', delimiter=',')
    # s2 = get_form(a2, b2, beta, v)
    s2 = get_form(a2, b2, beta, v)
    # s1 = get_form(a1, b1, xita, v)

    return s1, s2

def get_s5(s4:np.ndarray):
    print(s4.shape)
    return np.array([l for l in s4 if (l[0] != l[2]) or (l[1] != l[3]) or (l[4] != l[6]) or (l[5] != l[7])])

    # for o in range(len(s4)): # 每次都要输入，这里1144是粗筛选过后s4的列数,
    #     if s4[o, 0] != s4[o, 2] or s4[o,1] != s4[o, 3] or s4[o, 4] != s4[o, 6] or s4[o, 5] != s4[o, 7]: # 有一些夹角为0，因为配对了至少一对同样的晶格矢量（即相当于MS中输入了ab和ab），这里把0度的pass，但是180的没有pass
    #         s5(k,:)=s4(o,:)

def cross(a:Vec, b:Vec):
    return Vec([
        a.raw[1]*b.raw[2]-a.raw[2]*b.raw[1],
        -(a.raw[0]*b.raw[2]-a.raw[2]*b.raw[0]),
        a.raw[0]*b.raw[1]-a.raw[1]*b.raw[0]
    ])

def get_s6(s5:np.ndarray):
    a1 = 3.776
    b1 = 9.486
    xita = 90
    N1 = 36

    a2 = 4.78069
    b2 = 4.78069
    beta = 120
    N2 = 7

    mis_ang = 3
    mis_rate = 0.05

    xit = xita * pi / 180.0
    beit = beta * pi / 180.0

    # 如果这里算完，筛选出来的结果还是过多，必须要从原子的个数角度再进行筛选
    aa1 = Vec([cos(xit / 2) * a1, sin(xit / 2) * a1, 0])
    # 每次都要输入，定义第一个材料，在没有变化晶格矢量时。注意：这里都是默认c方向垂直于ab面
    bb1 = Vec([cos(xit / 2) * b1, sin(xit / 2) * b1 * (-1), 0])    # 每次都要输入
    S1 = norm(cross(aa1, bb1))
    # 叉乘法直接算第一个表面模型的面积（晶格矢量没有变化时）

    aa2 = Vec([cos(beit / 2) * a2, sin(beit / 2) * a2, 0])
    # 每次都要输入，定义第二个材料，在没有变化晶格矢量时。注意：这里都是默认c方向垂直于ab面
    bb2 = Vec([cos(beit / 2) * b2, sin(beit / 2) * b2 * (-1), 0])    # 每次都要输入
    S2 = norm(cross(aa2, bb2))
    # 法直接算第二个表面模型的面积（晶格矢量没有变化时）

    temp = []
    for i in s5:    # 每次都要输入，988为s5的列数
        m1, n1, m2, n2 = i[:4]
        AA1 = aa1 * m1 + bb1 * n1
        BB1 = aa1 * m2 + bb1 * n2
        S1_v1 = cross(AA1, BB1)
        S1_v = norm(S1_v1)
        N1_v = S1_v / S1 * N1
        # s5(i, 16) = N1_v    # 材料1改变晶格矢量后原子的个数


        j1, k1, j2, k2 = i[4: 8]
        AA2 = aa2 * j1 + bb2 * k1
        BB2 = aa2 * j2 + bb2 * k2
        S2_v1 = cross(AA2, BB2)
        S2_v = norm(S2_v1)
        N2_v = S2_v / S2 * N2
        # s5(i, 17) = N2_v; % 材料2改变晶格矢量后原子的个数
        # s5(i, 18) = N1_v + N2_v

        temp.append([N1_v, N2_v, N1_v+N2_v])
    temp = np.array(temp)
    s6 = np.hstack((s5, temp))
    s6 = s6[s6[:, 17].argsort()[::-1]]
    return s6


def get_real(x:complex):
    return x.real


def get_s7(s6):
    # 最后的筛选，新晶格矢量的配对是纯数学的两两组合，会遇到组合后两两平行，此时原子数为0，没有物理意义，将其剔除
    return np.array([list(map(get_real, o)) for o in s6 if o[15] and o[17]])
    # for o in s6:
    #     if o[15] and o[17]:
    #         s7(k,:)=s6(o,:);




if __name__ == '__main__':
    s1, s2 = main()

    pertime = time.perf_counter()
    with Pool() as p:
        result = np.array([p.apply_async(func=get_s3, args=(s1, s2, m, n)).get() for m in range(len(s1)) for n in range(len(s2))])

    np.savetxt('res.csv', result, fmt='%.5f', delimiter=',')
    print(result)
    print(result.shape)
    print(time.perf_counter() - pertime)
    s4 = get_s4(result)
    s5 = get_s5(s4)
    s6 = get_s6(s5)
    s7 = get_s7(s6)
    np.savetxt('s7.csv', s7, fmt='%.5f', delimiter=',')